import re
import textwrap

# imports
import pytest  # used for our unit tests
from codeflash.languages.java.instrumentation import \
    _add_timing_instrumentation

def test_no_test_annotations_returns_identical_source():
    # Basic: When there are no @Test annotations, source should be unchanged.
    src = textwrap.dedent(
        """\
        public class Example {
            public void helper() {
                // nothing to do
            }
        }"""
    )
    codeflash_output = _add_timing_instrumentation(src, "Example", "helper"); out = codeflash_output # 3.14μs -> 3.06μs (2.65% faster)

def test_single_test_method_inserts_instrumentation_and_preserves_body_order():
    # Basic: One @Test method with a simple single-line body.
    src = textwrap.dedent(
        """\
        public class MyClass {
            @Test
            public void testSomething() {
                System.out.println("hello");
            }
        }"""
    )
    codeflash_output = _add_timing_instrumentation(src, "MyClass", "testSomething"); out = codeflash_output # 11.9μs -> 10.3μs (15.1% faster)

    # The original body line should still be present (once) and should appear after the start println
    start_pos = out.find("!$######")
    body_pos = out.find('System.out.println("hello");')

def test_multiple_annotations_and_multiline_signature_preserved_and_instrumented():
    # Edge: Additional annotations should be preserved and multi-line method signature handled.
    src = textwrap.dedent(
        """\
        public class Multi {
            @Test
            @CustomAnnotation(value = "x")
            public void
            testMultiLine()
            throws Exception
            {
                // inner action
            }
        }"""
    )
    codeflash_output = _add_timing_instrumentation(src, "Multi", "testMultiLine"); out = codeflash_output # 12.3μs -> 11.0μs (12.2% faster)

def test_nested_braces_in_method_body_are_handled_correctly():
    # Edge: Method body contains nested braces, ensure correct brace counting and that nested content appears
    src = textwrap.dedent(
        """\
        public class Nested {
            @Test
            public void testNested() {
                for (int i = 0; i < 2; i++) {
                    if (i % 2 == 0) {
                        System.out.println(i);
                    }
                }
            }
        }"""
    )
    codeflash_output = _add_timing_instrumentation(src, "Nested", "testNested"); out = codeflash_output # 15.3μs -> 11.6μs (32.2% faster)

    # Confirm that the finally block for iter_id 1 exists and follows the try
    idx_try = out.find("try {")
    idx_finally = out.find("} finally {")

def test_missing_opening_brace_results_in_timing_start_inserted_but_no_finally():
    # Edge: @Test present but no opening '{' for the method in the source.
    src = textwrap.dedent(
        """\
        public class Incomplete {
            @Test
            public void testNoBrace()
        }"""
    )
    codeflash_output = _add_timing_instrumentation(src, "Incomplete", "testNoBrace"); out = codeflash_output # 7.46μs -> 7.05μs (5.83% faster)

def test_large_scale_many_tests_unique_iter_ids_and_counts():
    # Large Scale: Generate a source with many (@Test) methods to validate scalability and uniqueness of iter ids.
    # Use 100 methods (keeps under 1000).
    num_methods = 100
    methods = []
    for i in range(1, num_methods + 1):
        methods.append(textwrap.dedent(
            f"""\
            @Test
            public void test{i}() {{
                // body {i}
            }}"""
        ))
    class_src = "public class Bulk {\n" + "\n".join(methods) + "\n}"
    codeflash_output = _add_timing_instrumentation(class_src, "Bulk", "testUnknown"); out = codeflash_output # 436μs -> 396μs (10.1% faster)

    # Verify that we have exactly num_methods occurrences of loop variable declarations with increasing indices
    # E.g., int _cf_loop1, int _cf_loop2, ..., int _cf_loop100
    for idx in range(1, num_methods + 1):
        pattern = f"int _cf_loop{idx} = Integer.parseInt(System.getenv(\"CODEFLASH_LOOP_INDEX\"));"

    # There should be exactly num_methods start marker println occurrences (one per method) and exactly num_methods end markers.
    # Start marker uses "######$!" at end; end marker uses "######!".
    count_start_markers = out.count("######$!")
    count_end_markers = out.count("######!")

    # Ensure ordering: for the first method, the loop declaration appears before the first start println
    first_loop_idx = out.find("int _cf_loop1")
    first_start_idx = out.find("!$######")

def test_instrumentation_preserves_additional_annotation_blocks_and_spacing():
    # Basic/Edge: Ensure that blocks of annotations are preserved and spacing is not lost between annotation and signature.
    src = textwrap.dedent(
        """\
        public class Spacing {
            @Test
            @Repeat(3)
            @Another
            public void spacedTest() {
                doSomething();
            }
        }"""
    )
    codeflash_output = _add_timing_instrumentation(src, "Spacing", "spacedTest"); out = codeflash_output # 12.5μs -> 11.1μs (13.2% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import os

import pytest
from codeflash.languages.java.instrumentation import \
    _add_timing_instrumentation

def test_basic_single_test_method_instrumentation():
    """Test that a single @Test method is properly instrumented with timing code."""
    source = """public class TestExample {
    @Test
    public void testSimple() {
        int x = 5;
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testSimple"); result = codeflash_output # 12.2μs -> 11.0μs (11.1% faster)

def test_basic_method_signature_preserved():
    """Test that the method signature is preserved after instrumentation."""
    source = """public class TestExample {
    @Test
    public void testSimple() {
        int x = 5;
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testSimple"); result = codeflash_output # 11.9μs -> 10.5μs (13.9% faster)

def test_basic_class_and_func_names_in_output():
    """Test that class and function names are properly embedded in timing markers."""
    source = """public class MyTest {
    @Test
    public void myFunc() {
        x = 1;
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "MyTest", "myFunc"); result = codeflash_output # 11.5μs -> 10.4μs (11.1% faster)

def test_basic_multiple_test_methods():
    """Test that multiple @Test methods are each instrumented independently."""
    source = """public class TestExample {
    @Test
    public void testFirst() {
        int x = 1;
    }
    
    @Test
    public void testSecond() {
        int y = 2;
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testMethod"); result = codeflash_output # 19.1μs -> 16.8μs (13.5% faster)

def test_basic_single_line_method_body():
    """Test instrumentation of a method with a single-line body."""
    source = """public class TestExample {
    @Test
    public void test() {
        assert true;
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "test"); result = codeflash_output # 11.2μs -> 9.95μs (12.7% faster)

def test_basic_multiline_method_body():
    """Test instrumentation of a method with multiple lines in the body."""
    source = """public class TestExample {
    @Test
    public void test() {
        int a = 1;
        int b = 2;
        int c = a + b;
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "test"); result = codeflash_output # 13.1μs -> 10.7μs (22.4% faster)

def test_edge_empty_method_body():
    """Test instrumentation of a method with an empty body."""
    source = """public class TestExample {
    @Test
    public void testEmpty() {
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testEmpty"); result = codeflash_output # 10.2μs -> 8.97μs (13.3% faster)

def test_edge_method_with_multiple_annotations():
    """Test instrumentation when method has multiple annotations."""
    source = """public class TestExample {
    @Test
    @Ignore
    @DisplayName("My Test")
    public void testMultiAnnot() {
        x = 1;
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testMultiAnnot"); result = codeflash_output # 12.0μs -> 10.6μs (13.4% faster)

def test_edge_test_annotation_with_parameters():
    """Test instrumentation when @Test has parameters."""
    source = """public class TestExample {
    @Test(timeout = 1000)
    public void testWithParams() {
        int x = 5;
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testWithParams"); result = codeflash_output # 11.3μs -> 9.80μs (15.4% faster)

def test_edge_nested_braces_in_method_body():
    """Test instrumentation when method body contains nested braces."""
    source = """public class TestExample {
    @Test
    public void testNested() {
        if (true) {
            int x = 1;
        }
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testNested"); result = codeflash_output # 13.0μs -> 10.9μs (19.1% faster)

def test_edge_special_characters_in_names():
    """Test instrumentation when class/function names contain underscores."""
    source = """public class Test_Example_Class {
    @Test
    public void test_Function_Name() {
        int x = 1;
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "Test_Example_Class", "test_Function_Name"); result = codeflash_output # 11.3μs -> 10.0μs (12.6% faster)

def test_edge_method_with_parameters():
    """Test instrumentation of a test method that has parameters."""
    source = """public class TestExample {
    @Test
    public void testWithParams(int param1, String param2) {
        assert param1 > 0;
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testWithParams"); result = codeflash_output # 11.4μs -> 9.84μs (16.1% faster)

def test_edge_method_with_throws_clause():
    """Test instrumentation of a test method with throws clause."""
    source = """public class TestExample {
    @Test
    public void testWithThrows() throws Exception {
        int x = 1;
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testWithThrows"); result = codeflash_output # 11.3μs -> 9.80μs (15.3% faster)

def test_edge_method_with_braces_in_string():
    """Test instrumentation when method body contains braces inside strings."""
    source = """public class TestExample {
    @Test
    public void testStringBraces() {
        String s = "test { braces }";
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testStringBraces"); result = codeflash_output # 11.9μs -> 9.86μs (20.9% faster)

def test_edge_indented_class():
    """Test instrumentation when the test method is indented within a class."""
    source = """    public class TestExample {
        @Test
        public void testIndented() {
            int x = 1;
        }
    }"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testIndented"); result = codeflash_output # 11.9μs -> 10.3μs (15.1% faster)

def test_edge_no_test_annotation():
    """Test that methods without @Test annotation are not instrumented."""
    source = """public class TestExample {
    public void normalMethod() {
        int x = 1;
    }
    
    @Test
    public void testMethod() {
        int y = 2;
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "test"); result = codeflash_output # 13.0μs -> 11.2μs (16.1% faster)

def test_edge_non_ascii_characters_in_body():
    """Test that non-ASCII characters in method body are preserved."""
    source = """public class TestExample {
    @Test
    public void testUnicode() {
        String s = "café";
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testUnicode"); result = codeflash_output # 12.7μs -> 10.8μs (17.9% faster)

def test_edge_very_long_method_name():
    """Test instrumentation with very long class and method names."""
    long_class = "VeryLongTestClassNameThatIsActuallyQuiteLong"
    long_func = "testWithAVeryLongFunctionNameThatExplainsWhatItDoes"
    
    source = f"""public class {long_class} {{
    @Test
    public void {long_func}() {{
        int x = 1;
    }}
}}"""
    
    codeflash_output = _add_timing_instrumentation(source, long_class, long_func); result = codeflash_output # 11.5μs -> 10.2μs (13.3% faster)

def test_edge_method_with_only_comment_in_body():
    """Test instrumentation when method body contains only comments."""
    source = """public class TestExample {
    @Test
    public void testOnlyComment() {
        // This is a comment
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testOnlyComment"); result = codeflash_output # 11.5μs -> 10.1μs (14.0% faster)

def test_edge_class_name_with_package():
    """Test instrumentation when class name contains package notation."""
    source = """public class TestExample {
    @Test
    public void testMethod() {
        int x = 1;
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "com.example.TestExample", "testMethod"); result = codeflash_output # 11.2μs -> 9.92μs (13.2% faster)

def test_large_scale_many_test_methods():
    """Test instrumentation of a class with many @Test methods."""
    # Create a source with 100 test methods
    lines = ["public class TestExample {"]
    for i in range(100):
        lines.append(f"    @Test")
        lines.append(f"    public void test{i}() {{")
        lines.append(f"        int x = {i};")
        lines.append(f"    }}")
        lines.append("")
    lines.append("}")
    
    source = "\n".join(lines)
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "test"); result = codeflash_output # 494μs -> 430μs (14.7% faster)

def test_large_scale_large_method_body():
    """Test instrumentation of a method with a large body."""
    # Create method with 200 lines in body
    lines = ["public class TestExample {", "    @Test", "    public void testLargeBody() {"]
    for i in range(200):
        lines.append(f"        int var{i} = {i};")
    lines.append("    }")
    lines.append("}")
    
    source = "\n".join(lines)
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testLargeBody"); result = codeflash_output # 173μs -> 76.6μs (127% faster)

def test_large_scale_deeply_nested_braces():
    """Test instrumentation with deeply nested brace structures."""
    # Create method with deeply nested braces
    lines = ["public class TestExample {", "    @Test", "    public void testDeepNesting() {"]
    indent_level = 2
    for i in range(20):
        lines.append("    " * indent_level + "if (true) {")
        indent_level += 1
    lines.append("    " * indent_level + "int x = 1;")
    for i in range(20):
        indent_level -= 1
        lines.append("    " * indent_level + "}")
    lines.append("    }")
    lines.append("}")
    
    source = "\n".join(lines)
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testDeepNesting"); result = codeflash_output # 76.6μs -> 26.2μs (192% faster)

def test_large_scale_mixed_methods():
    """Test instrumentation with a mix of test and non-test methods."""
    # Create class with 50 @Test methods and 50 non-test methods
    lines = ["public class TestExample {"]
    for i in range(50):
        lines.append(f"    @Test")
        lines.append(f"    public void test{i}() {{")
        lines.append(f"        int x = {i};")
        lines.append(f"    }}")
        lines.append(f"    public void helper{i}() {{")
        lines.append(f"        int y = {i};")
        lines.append(f"    }}")
        lines.append("")
    lines.append("}")
    
    source = "\n".join(lines)
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "test"); result = codeflash_output # 284μs -> 249μs (14.0% faster)
    # Non-test helper methods should remain unchanged
    for i in range(50):
        pass

def test_large_scale_very_long_source_file():
    """Test instrumentation on a very large source file."""
    # Create source with 500 total lines (mix of test and other code)
    lines = [
        "package com.example;",
        "import java.util.*;",
        "public class LargeTest {"
    ]
    
    # Add various class members and 50 test methods
    for i in range(50):
        lines.append(f"    private static final int CONSTANT_{i} = {i};")
    
    lines.append("")
    
    for i in range(50):
        lines.append(f"    @Test")
        lines.append(f"    public void testMethod{i}() {{")
        lines.append(f"        assertEquals({i}, CONSTANT_{i});")
        lines.append(f"    }}")
        lines.append("")
    
    lines.append("}")
    source = "\n".join(lines)
    codeflash_output = _add_timing_instrumentation(source, "LargeTest", "test"); result = codeflash_output # 285μs -> 228μs (24.5% faster)

def test_large_scale_complex_method_bodies():
    """Test instrumentation of methods with complex control flow."""
    source = """public class TestExample {
    @Test
    public void testComplex() {
        for (int i = 0; i < 100; i++) {
            if (i % 2 == 0) {
                switch (i) {
                    case 0:
                        int x = 0;
                        break;
                    case 2:
                        int y = 2;
                        break;
                    default:
                        int z = i;
                }
            }
        }
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testComplex"); result = codeflash_output # 24.7μs -> 15.9μs (55.6% faster)

def test_large_scale_variable_naming_no_conflicts():
    """Test that generated variable names don't conflict with user code."""
    # Create method that uses variables with similar names to generated ones
    source = """public class TestExample {
    @Test
    public void testNoConflict() {
        int _cf_loop = 5;
        int _cf_i = 10;
        int _cf_start = 20;
        for (int i = 0; i < 10; i++) {
            int x = _cf_loop + _cf_i + _cf_start;
        }
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testNoConflict"); result = codeflash_output # 17.2μs -> 12.5μs (37.7% faster)

def test_large_scale_stress_output_length():
    """Test that output size remains reasonable with many instrumented methods."""
    # Create 200 test methods to stress output size
    lines = ["public class StressTest {"]
    for i in range(200):
        lines.append(f"    @Test")
        lines.append(f"    public void test{i}() {{")
        lines.append(f"        assertTrue(true);")
        lines.append(f"    }}")
    lines.append("}")
    
    source = "\n".join(lines)
    codeflash_output = _add_timing_instrumentation(source, "StressTest", "test"); result = codeflash_output # 1.03ms -> 873μs (17.9% faster)

def test_large_scale_whitespace_preservation():
    """Test that significant whitespace variations are handled correctly."""
    # Create source with inconsistent indentation and spacing
    source = """public class TestExample {
	@Test
	public void testWhitespace() {
  int x=1;
    int     y  =  2;
			int z = 3;
	}
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testWhitespace"); result = codeflash_output # 13.2μs -> 11.6μs (13.8% faster)

def test_integration_complete_class_with_multiple_features():
    """Test a complete realistic test class with various features."""
    source = """public class CompleteTest {
    private static final String TEST_DATA = "data";
    
    @Before
    public void setUp() {
        // Setup
    }
    
    @Test
    public void testFirst() {
        String result = processData(TEST_DATA);
        assertEquals("expected", result);
    }
    
    @Test
    public void testSecond() {
        try {
            doSomething();
        } catch (Exception e) {
            fail("Should not throw");
        }
    }
    
    @After
    public void tearDown() {
        // Cleanup
    }
    
    private String processData(String input) {
        return input.toUpperCase();
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "CompleteTest", "test"); result = codeflash_output # 29.8μs -> 23.8μs (25.1% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.